﻿#include "Engine.h"
#include <boost/chrono.hpp>
#include "../util/TextUtil.h"
#include "../util/DateTime.h"
#include "../log/SysLog.h"
#include "../service/ModelService.h"
#include "../service/SceneService.h"

Engine::Engine()
	: m_run(true)
	, m_avg_time(0.0001)
{
}

Engine::~Engine()
{
}

void Engine::Start()
{
	m_run = true;
	m_ocr_thread = std::thread(std::bind(&Engine::_Recognize, shared_from_this()));
	m_sync_thread = std::thread(std::bind(&Engine::_Synchronize, shared_from_this()));
}

void Engine::Stop()
{
	m_run = false;
	if (m_sync_thread.joinable())
	{
		m_sync_thread.join();
	}
	if (m_ocr_thread.joinable())
	{
		m_ocr_thread.join();
	}
}

void Engine::Push(OcrItemPtr item)
{
	// 1秒动态识别张数
	int image_count_per_sec = 1 / m_avg_time;
	// 超10秒才能识别完
	if (m_queue.size() > 10 * image_count_per_sec)
	{
		//LOG_WARN("identification quantity [{}] exceeds device capacity [avg:{:.3f}], ignore it.", m_queue.size(), m_avg_time);
		return;
	}
	// 能力足够则识别
	std::unique_lock<std::mutex> _guard(m_ocr_lock);
	m_queue.push(item);
}

OcrItemPtr Engine::Pop()
{
	OcrItemPtr item;
	std::unique_lock<std::mutex> _guard(m_ocr_lock);
	if (m_queue.size() > 0)
	{
		item = m_queue.front();
		m_queue.pop();
	}
	return item;
}

bool Engine::LoadModel(std::string id)
{
	ModelServicePtr modelService = std::make_shared<ModelService>();
	if (modelService)
	{
		ModelInfo model;
		model.SetId(id);
		if (modelService->Find(model))
		{
			std::unique_lock<std::mutex> _guard(m_infer_lock);
			if (m_infer_map.find(model.GetId()) == m_infer_map.end())
			{
				cv::Size image_size = cv::Size(model.GetWidth(), model.GetHeight());
				std::vector<std::string> classes = TextUtil::get_mutable_instance().SplitString(model.GetCls(), ',');
				InferencePtr inference = std::make_shared<Inference>(model, classes, image_size, true);
				if (!inference->loadOnnxNetwork())
				{
					LOG_INFO("载入onnx模型:{} 失败", model.GetId());
					return true;
				}
				else
				{
					m_infer_map.insert(std::make_pair(model.GetId(), inference));
					LOG_INFO("载入onnx模型:{} 成功", model.GetId());
				}
			}
		}
	}
	return false;
}

void Engine::SyncModel()
{
	// 查询系统模型
	std::vector<ModelInfo> models;
	ModelServicePtr modelService = std::make_shared<ModelService>();
	int pageIndex = 0, pageSize = 10;
	while (true)
	{
		std::vector<ModelInfo> list;
		modelService->FindPage(pageIndex, pageSize, "1", list);
		if (list.size() <= 0)
		{
			break;
		}
		models.insert(models.end(), list.begin(), list.end());
		pageIndex += pageSize;
	}

	// 添加所有模型
	for (size_t index = 0; index < models.size(); ++index)
	{
		ModelInfo& model = models[index];
		std::unique_lock<std::mutex> _guard(m_infer_lock);
		if (m_infer_map.find(model.GetId()) == m_infer_map.end())
		{
			cv::Size image_size = cv::Size(model.GetWidth(), model.GetHeight());;
			std::vector<std::string> classes = TextUtil::get_mutable_instance().SplitString(model.GetCls(), ',');
			InferencePtr inference = std::make_shared<Inference>(model, classes, image_size, true);
			if (!inference->loadOnnxNetwork())
			{
				LOG_WARN("load onnx model:{} failed.", model.GetId());
			}
			else
			{
				m_infer_map.insert(std::make_pair(model.GetId(), inference));
				LOG_INFO("load onnx model:{} success.", model.GetId());
			}
		}
	}

	// 移除不存在模型
	std::unique_lock<std::mutex> _guard(m_infer_lock);
	for (INFER_MAP::iterator it = m_infer_map.begin(); it != m_infer_map.end(); )
	{
		// 模型查询
		bool find = false;
		for (size_t index = 0; index < models.size(); ++index)
		{
			ModelInfo& model = models[index];
			if (model.GetId() == it->first)
			{
				// 找到模型且发生变化
				if (0 == DateTime::Diff(model.GetCreateTime(), it->second->GetModel().GetCreateTime()))
				{
					find = true;
					break;
				}
			}
		}
		// 模型已删除
		if (!find)
		{
			it = m_infer_map.erase(it);
		}
		// 模型未删除
		else
		{
			++it;
		}
	}
}

void Engine::SyncScene()
{
	// 查询所有场景
	std::vector<SceneInfo> scenes;
	SceneServicePtr sceneService = std::make_shared<SceneService>();
	int pageIndex = 0, pageSize = 10;
	while (true)
	{
		std::vector<SceneInfo> list;
		sceneService->FindPage(pageIndex, pageSize, list);
		if (list.size() <= 0)
		{
			break;
		}
		scenes.insert(scenes.end(), list.begin(), list.end());
		pageIndex += pageSize;
	}

	// 添加所有场景
	for (size_t index = 0; index < scenes.size(); ++index)
	{
		SceneInfo& scene = scenes[index];

		// 不存在设备场景
		std::unique_lock<std::mutex> _guard(m_scene_lock);
		if (m_device_scenes.find(scene.GetDevice()) == m_device_scenes.end())
		{
			SCENES dev_scenes = { scene };
			m_device_scenes[scene.GetDevice()] = dev_scenes;
			continue;
		}

		// 查找场景列表
		bool find = false;
		SCENES& dev_scenes = m_device_scenes[scene.GetDevice()];
		for (SCENES::iterator it = dev_scenes.begin(); it != dev_scenes.end(); ++it)
		{
			if (it->GetId() == scene.GetId())
			{
				find = true;
				break;
			}
		}

		// 不存在则新增
		if (!find)
		{
			dev_scenes.push_back(scene);
		}
	}

	// 移除不存在场景
	std::unique_lock<std::mutex> _guard(m_scene_lock);
	for (DEV_SCENES::iterator it = m_device_scenes.begin(); it != m_device_scenes.end(); ++it)
	{
		// 内存设备场景
		SCENES& dev_scenes = it->second;
		for (SCENES::iterator sit = dev_scenes.begin(); sit != dev_scenes.end();)
		{
			// 数据库是否存在
			bool find = false;
			for (size_t index = 0; index < scenes.size(); ++index)
			{
				// 找到且场景未变化
				SceneInfo& scene = scenes[index];
				if (scene.GetId() == sit->GetId() &&
					0 == DateTime::Diff(scene.GetUpdateTime(), sit->GetUpdateTime()))
				{
					find = true;
					break;
				}
			}
			// 场景已删或更新
			if (!find)
			{
				sit = dev_scenes.erase(sit);
			}
			else
			{
				++sit;
			}
		}
	}
}

void Engine::_Synchronize()
{
	// 数据同步
	while (m_run)
	{
		// 同步模型
		this->SyncModel();
		// 同步场景
		this->SyncScene();
		// 每10秒一次
		for (int index = 0; index < 10; ++index)
		{
			if (!m_run)
			{
				return;
			}
			std::this_thread::sleep_for(std::chrono::seconds(1));
		}
	}
}

void Engine::_Recognize()
{
	// 模型检测
	while (m_run)
	{
		// 获取识别项
		OcrItemPtr item = this->Pop();
		if (!item)
		{
			std::this_thread::sleep_for(std::chrono::milliseconds(10));
			continue;
		}

		// 获取分析场景
		SCENES scenes;
		{
			// 动态载入场景
			std::unique_lock<std::mutex> _guard(m_scene_lock);
			if (m_device_scenes.find(item->GetDevice()) == m_device_scenes.end())
			{
				continue;
			}
			// 获取场景信息
			scenes = m_device_scenes[item->GetDevice()];
		}

		// 获取模型推理
		std::vector<SCENE_INFE> inferences;
		{
			// 获取所有场景模型
			std::unique_lock<std::mutex> _guard(m_infer_lock);
			for (SCENES::iterator it = scenes.begin(); it != scenes.end(); ++it)
			{
				if (m_infer_map.find(it->GetModel()) == m_infer_map.end())
				{
					LOG_WARN("not found device:{} scene:{} model:{}, ignore it.", it->GetDevice(), it->GetId(), it->GetModel());
					continue;
				}
				InferencePtr inference = m_infer_map[it->GetModel()];
				SCENE_INFE scene_infe = { it->GetId(), *it, inference };
				inferences.push_back(scene_infe);
			}
		}
		if (inferences.size() <= 0)
		{
			LOG_WARN("not found device:{} model, ignore it.", item->GetDevice());
			continue;
		}

		// 获取开始时间点
		boost::chrono::system_clock::time_point start = DateTime::CurrrentTimePoint();
		// 当前图片推理
		size_t count = _Inference(inferences, item);
		// 获取结束时间
		double use_time = DateTime::Duration(start);
		m_times.push(use_time);
		// 最近10张平均时长
		if (m_times.size() > 10)
		{
			m_times.pop();
		}
		double avg = 0;
		std::queue<double> temp(m_times);
		while (!temp.empty()) {
			avg += temp.front();
			temp.pop();
		}
		m_avg_time = avg / m_times.size();
		if (count > 0)
		{
			LOG_INFO("recognize image: {}, avg:{}", use_time, m_avg_time);
		}
	}

	LOG_INFO("recognize engine exit.");
}

size_t Engine::_Inference(SCENE_INFES infes, OcrItemPtr item)
{
	// 打印推理结果
	std::vector<OCR_RESULT> results;

	// 按场景逐一推理
	for (SCENE_INFES::iterator it = infes.begin(); it != infes.end(); ++it)
	{
		// 场景模型信息
		SCENE_INFE scene_infer = *it;
		SceneInfo& scene = std::get<1>(scene_infer);

		// 异常分类表
		std::vector<std::string> alarms = TextUtil::get_mutable_instance().SplitString(scene.GetDetectType(), ',');

		// 推理模型无效
		InferencePtr inference = std::get<2>(scene_infer);
		if (!inference)
		{
			continue;
		}

		// 对图片进行推理
		std::vector<Detection> detections = inference->runInference(item->GetImage());
		if (detections.size() <= 0)
		{
			continue;
		}

		// 提取分析结果
		std::map<std::string, OCR_RESULT> class_result;
		for (size_t i = 0; i < detections.size(); ++i)
		{
			// 获取结果
			Detection detection = detections[i];
			// 边框颜色
			cv::Rect box = detection.box;
			cv::Scalar color = detection.color;
			// 报警匹配
			if (TextUtil::get_mutable_instance().ContainString(alarms, detection.className))
			{
				// BGR-红色
				color = cv::Scalar(0, 0, 255);

				// 精度不符
				int confidence = detection.confidence * 100;
				if (confidence < scene.GetThreshold())
				{
					continue;
				}

				// 识别结果
				OCR_RESULT result;
				result.device = scene.GetDevice();
				result.alarmType = scene.GetAlarmCode();
				result.alarmInfo = scene.GetAlarmInfo();
				result.confidence = confidence;
				result.count = 1;

				// 保存最大精度
				if (class_result.find(detection.className) != class_result.end())
				{
					result.count = class_result[detection.className].count + 1;
					if (confidence > class_result[detection.className].confidence)
					{
						class_result[detection.className] = result;
					}
				}
				else
				{
					class_result[detection.className] = result;
				}
			}
			else
			{
				// BGR-绿色
				color = cv::Scalar(0, 255, 0);
			}

			// 绘制边框
			cv::rectangle(item->GetImage(), box, color, 2);
			// 分类名+置信度
			std::string classString = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
			cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
			// 标签显示位置
			cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);
			// 绘制标签盒子
			cv::rectangle(item->GetImage(), textBox, color, cv::FILLED);
			// 绘制标签文本
			cv::putText(item->GetImage(), classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
		}

		// 汇总分析结果
		for (std::map<std::string, OCR_RESULT>::iterator it = class_result.begin(); it != class_result.end(); ++it)
		{
			results.push_back(it->second);
		}
	}

	// 未分析到报警
	if (results.size() <= 0)
	{
		return 0;
	}

	// 报警信息回调
	if (item->GetCallback())
	{
		item->GetCallback()(item->GetDevice(), results, item->GetImage(), item->GetContext());
	}

	return results.size();
}